#include <mpi.h>

#include <cuda_runtime.h>

#include <stdio.h>

#define min(x, y) ( ((x) < (y)) ? (x) : (y) )

void Kernel_Wrapper(float *, int, int);

int main(int argc, char * argv[])
{
	MPI_Init(&argc, &argv);

	int rank, size;

	MPI_Comm_rank(MPI_COMM_WORLD, &rank);

	MPI_Comm_size(MPI_COMM_WORLD, &size);

	int device_count = 0;

	cudaGetDeviceCount(&device_count);

	int set_device = min(device_count-1, rank % size);

	cudaSetDevice(set_device);

	fprintf(stdout, "Rank %d size %d device %d / %d\n", rank, size, set_device, device_count);

	float * dData = NULL;
	float * hData = NULL;

	int N = 32;

	cudaMallocHost((void **) &hData, N * sizeof(float));
	cudaMalloc((void **) &dData, N * sizeof(float));

	Kernel_Wrapper(dData, rank, N);

	cudaMemcpy(hData, dData, N * sizeof(float), cudaMemcpyDeviceToHost);

	fprintf(stdout, "0 : %f | 1 : %f | 2 : %f | 3 : %f\n", hData[0], hData[1], hData[2], hData[3]);

	float recvbuf = 0.0f;
	MPI_Status st;

	MPI_Send(&hData[3], 1, MPI_FLOAT, (rank+1) % size, 0, MPI_COMM_WORLD);
	MPI_Recv(&recvbuf, 1, MPI_FLOAT, (size + rank-1) % size, 0, MPI_COMM_WORLD, &st);

	fprintf(stdout, "Rank %d rec %f\n", rank, recvbuf);

	cudaFree(dData);
	cudaFreeHost(hData);

	MPI_Finalize();

	return 0;
}